"""
Author:cold
Date:2021-04-04
Version:3.0
Info:baseline
"""
from pandas import read_csv
from sklearn.model_selection import train_test_split
from sklearn.linear_model import LinearRegression
from sklearn.metrics import mean_squared_error
import pandas as pd
from sklearn.preprocessing import StandardScaler
from sklearn.feature_selection import SelectKBest,f_regression
from matplotlib import pyplot as plt


#字典结果：{提取k个最佳特征，及索引}
def ToBeStdAndSel(X,Y,k):

    best = SelectKBest(f_regression, k=k)
    X_best = best.fit_transform(X, Y)#A
    best_index = best.get_support()#B AB顺序不能换
    BEST = {}
    BEST['best_index'] = best_index
    BEST['X_best'] = X_best # 等价于 X_std[:,best_index]
    return BEST

#标准化
def ToBeStd(X):
    stand = StandardScaler()
    X_std = stand.fit_transform(X)
    return X_std

# 保存csv
def ToSaveCsv(y_pred):
    ID = []
    for i in range(len(y_pred)):
        ID.append("id_" + str(i + 1))
    res = pd.DataFrame()
    res['ID'] = ID
    res['value'] = y_pred
    res.to_csv('res.csv', index=False)
    print("res.csv 已生成")

#预测
def TryToPredict(testset):
    x_pred = testset[:, 1:14]
    x_pred_best = x_pred[:, best_index]
    y_pred = lr.predict(x_pred_best)
    return y_pred


# 加载数据(455)
dataset =read_csv('train_dataset.csv').values
# 划分训练集和测试集（+数据标准化,+特征提取）X--> X_std--> X_best
X = dataset[:,0:13]
Y = dataset[:,13]
BEST = ToBeStdAndSel(X,Y,6)
X_best = BEST['X_best']
best_index = BEST['best_index']

x_train,x_test,y_train,y_test = train_test_split(X_best,Y,test_size=0.3)

# 创建线性回归模型
lr = LinearRegression()
# 拟合训练数据
lr.fit(x_train,y_train)
# 得到预测结果
y_test_pred = lr.predict(x_test)
y_train_pred = lr.predict(x_train)


# 计算相应的评测指标
error_test = mean_squared_error(y_test,y_test_pred)
error_train = mean_squared_error(y_train,y_train_pred)
print("训练集误差为：{}，测试集误差为：{}".format(error_train,error_test))
plt.plot(y_test_pred,'r-',label='predict_value')
plt.plot(y_test,'b-',label='true_value')
plt.legend()
plt.show()

#预测、保存
testset =read_csv('test_dataset.csv').values
y_pred = TryToPredict(testset)
ToSaveCsv(y_pred)

